#include <vtkSmartPointer.h>
#include <vtkUnstructuredGrid.h>
#include <vtkPoints.h>
#include <vtkCellArray.h>
#include <vtkTetra.h>
#include <vtkHexahedron.h>
#include <vtkDataSetMapper.h>
#include <vtkActor.h>
#include <vtkRenderer.h>
#include <vtkRenderWindow.h>
#include <vtkRenderWindowInteractor.h>
#include <vtkInteractorStyleTrackballCamera.h>
#include <vtkCellPicker.h>
#include <vtkProperty.h>
#include <vtkCommand.h>
#include <vtkIdTypeArray.h>
#include <vtkSelectionNode.h>
#include <vtkSelection.h>
#include <vtkExtractSelection.h>
#include <vtkDataSetMapper.h>
#include <vtkLookupTable.h>
#include <vtkUnsignedCharArray.h>
#include <vtkPointData.h>
#include <vtkCellData.h>
#include <vtkRendererCollection.h>

// 自定义交互器样式类，用于处理鼠标点击事件
class MouseInteractorStyle : public vtkInteractorStyleTrackballCamera
{
public:
    static MouseInteractorStyle* New();
    vtkTypeMacro(MouseInteractorStyle, vtkInteractorStyleTrackballCamera);

    virtual void OnLeftButtonDown() override
    {
        // 获取鼠标点击位置
        int* clickPos = this->GetInteractor()->GetEventPosition();

        // 使用CellPicker进行拾取
        vtkSmartPointer<vtkCellPicker> picker = vtkSmartPointer<vtkCellPicker>::New();
        this->GetInteractor()->SetPicker(picker);
        picker->SetTolerance(0.0005);
        
        // 执行拾取操作
        picker->Pick(clickPos[0], clickPos[1], 0, this->GetInteractor()->GetRenderWindow()->GetRenderers()->GetFirstRenderer());
        
        // 获取拾取到的cell ID
        vtkIdType cellId = picker->GetCellId();
        
        if (cellId != -1)
        {
            std::cout << "拾取到Cell ID: " << cellId << std::endl;
            
            // 高亮被选中的cell
            this->HighlightPickedCell(cellId);
        }
        else
        {
            std::cout << "未拾取到任何cell" << std::endl;
            // 清除高亮
            this->ClearHighlight();
        }

        // 调用父类的方法
        vtkInteractorStyleTrackballCamera::OnLeftButtonDown();
    }

    void SetData(vtkUnstructuredGrid* data)
    {
        this->Data = data;
        this->OriginalMapper = vtkSmartPointer<vtkDataSetMapper>::New();
        this->OriginalMapper->SetInputData(data);
        
        this->HighlightMapper = vtkSmartPointer<vtkDataSetMapper>::New();
        
        // 为原始数据创建颜色数组
        this->CreateOriginalColors();
    }

    void SetActor(vtkActor* actor)
    {
        this->Actor = actor;
    }

private:
    vtkUnstructuredGrid* Data;
    vtkSmartPointer<vtkDataSetMapper> OriginalMapper;
    vtkSmartPointer<vtkDataSetMapper> HighlightMapper;
    vtkActor* Actor;
    vtkIdType LastPickedCell = -1;

    void CreateOriginalColors()
    {
        // 创建颜色数组，给所有cell默认颜色（蓝色）
        vtkSmartPointer<vtkUnsignedCharArray> colors = vtkSmartPointer<vtkUnsignedCharArray>::New();
        colors->SetNumberOfComponents(3);
        colors->SetName("Colors");
        
        for (vtkIdType i = 0; i < this->Data->GetNumberOfCells(); i++)
        {
            unsigned char color[3] = {100, 149, 237}; // 蓝色
            colors->InsertNextTypedTuple(color);
        }
        
        this->Data->GetCellData()->SetScalars(colors);
    }

    void HighlightPickedCell(vtkIdType cellId)
    {
        if (this->Data)
        {
            // 重置所有颜色为默认颜色
            vtkUnsignedCharArray* colors = vtkUnsignedCharArray::SafeDownCast(
                this->Data->GetCellData()->GetScalars("Colors"));
            
            if (colors)
            {
                // 重置所有cell为蓝色
                for (vtkIdType i = 0; i < this->Data->GetNumberOfCells(); i++)
                {
                    unsigned char color[3] = {100, 149, 237}; // 蓝色
                    colors->SetTypedTuple(i, color);
                }
                
                // 设置被选中的cell为红色
                unsigned char highlightColor[3] = {255, 0, 0}; // 红色
                colors->SetTypedTuple(cellId, highlightColor);
                
                // 更新数据
                colors->Modified();
                this->Data->Modified();
                
                // 重新渲染
                this->GetInteractor()->GetRenderWindow()->Render();
                
                this->LastPickedCell = cellId;
            }
        }
    }

    void ClearHighlight()
    {
        if (this->Data && this->LastPickedCell != -1)
        {
            // 将之前高亮的cell恢复为默认颜色
            vtkUnsignedCharArray* colors = vtkUnsignedCharArray::SafeDownCast(
                this->Data->GetCellData()->GetScalars("Colors"));
            
            if (colors)
            {
                unsigned char defaultColor[3] = {100, 149, 237}; // 蓝色
                colors->SetTypedTuple(this->LastPickedCell, defaultColor);
                colors->Modified();
                this->Data->Modified();
                this->GetInteractor()->GetRenderWindow()->Render();
            }
            
            this->LastPickedCell = -1;
        }
    }
};

vtkStandardNewMacro(MouseInteractorStyle);

// 创建示例非结构化网格的函数
vtkSmartPointer<vtkUnstructuredGrid> CreateSampleUnstructuredGrid()
{
    vtkSmartPointer<vtkUnstructuredGrid> ugrid = vtkSmartPointer<vtkUnstructuredGrid>::New();
    
    // 创建点
    vtkSmartPointer<vtkPoints> points = vtkSmartPointer<vtkPoints>::New();
    
    // 创建一个包含多个四面体和六面体的网格
    // 四面体的点
    points->InsertNextPoint(0.0, 0.0, 0.0);  // 0
    points->InsertNextPoint(1.0, 0.0, 0.0);  // 1
    points->InsertNextPoint(0.5, 1.0, 0.0);  // 2
    points->InsertNextPoint(0.5, 0.5, 1.0);  // 3
    
    // 六面体的点
    points->InsertNextPoint(2.0, 0.0, 0.0);  // 4
    points->InsertNextPoint(3.0, 0.0, 0.0);  // 5
    points->InsertNextPoint(3.0, 1.0, 0.0);  // 6
    points->InsertNextPoint(2.0, 1.0, 0.0);  // 7
    points->InsertNextPoint(2.0, 0.0, 1.0);  // 8
    points->InsertNextPoint(3.0, 0.0, 1.0);  // 9
    points->InsertNextPoint(3.0, 1.0, 1.0);  // 10
    points->InsertNextPoint(2.0, 1.0, 1.0);  // 11
    
    // 第二个四面体的点
    points->InsertNextPoint(0.0, 2.0, 0.0);  // 12
    points->InsertNextPoint(1.0, 2.0, 0.0);  // 13
    points->InsertNextPoint(0.5, 3.0, 0.0);  // 14
    points->InsertNextPoint(0.5, 2.5, 1.0);  // 15
    
    ugrid->SetPoints(points);
    
    // 创建第一个四面体
    vtkSmartPointer<vtkTetra> tetra1 = vtkSmartPointer<vtkTetra>::New();
    tetra1->GetPointIds()->SetId(0, 0);
    tetra1->GetPointIds()->SetId(1, 1);
    tetra1->GetPointIds()->SetId(2, 2);
    tetra1->GetPointIds()->SetId(3, 3);
    ugrid->InsertNextCell(tetra1->GetCellType(), tetra1->GetPointIds());
    
    // 创建六面体
    vtkSmartPointer<vtkHexahedron> hex = vtkSmartPointer<vtkHexahedron>::New();
    hex->GetPointIds()->SetId(0, 4);
    hex->GetPointIds()->SetId(1, 5);
    hex->GetPointIds()->SetId(2, 6);
    hex->GetPointIds()->SetId(3, 7);
    hex->GetPointIds()->SetId(4, 8);
    hex->GetPointIds()->SetId(5, 9);
    hex->GetPointIds()->SetId(6, 10);
    hex->GetPointIds()->SetId(7, 11);
    ugrid->InsertNextCell(hex->GetCellType(), hex->GetPointIds());
    
    // 创建第二个四面体
    vtkSmartPointer<vtkTetra> tetra2 = vtkSmartPointer<vtkTetra>::New();
    tetra2->GetPointIds()->SetId(0, 12);
    tetra2->GetPointIds()->SetId(1, 13);
    tetra2->GetPointIds()->SetId(2, 14);
    tetra2->GetPointIds()->SetId(3, 15);
    ugrid->InsertNextCell(tetra2->GetCellType(), tetra2->GetPointIds());
    
    return ugrid;
}

int main()
{
    system("chcp 65001");
    // 创建非结构化网格
    vtkSmartPointer<vtkUnstructuredGrid> ugrid = CreateSampleUnstructuredGrid();
    
    // 创建映射器
    vtkSmartPointer<vtkDataSetMapper> mapper = vtkSmartPointer<vtkDataSetMapper>::New();
    mapper->SetInputData(ugrid);
    
    // 创建演员
    vtkSmartPointer<vtkActor> actor = vtkSmartPointer<vtkActor>::New();
    actor->SetMapper(mapper);
    actor->GetProperty()->SetRepresentationToSurface();
    actor->GetProperty()->EdgeVisibilityOn(); // 显示边框
    actor->GetProperty()->SetEdgeColor(0.0, 0.0, 0.0); // 黑色边框
    
    // 创建渲染器
    vtkSmartPointer<vtkRenderer> renderer = vtkSmartPointer<vtkRenderer>::New();
    renderer->AddActor(actor);
    renderer->SetBackground(0.2, 0.3, 0.4); // 设置背景颜色
    
    // 创建渲染窗口
    vtkSmartPointer<vtkRenderWindow> renderWindow = vtkSmartPointer<vtkRenderWindow>::New();
    renderWindow->AddRenderer(renderer);
    renderWindow->SetWindowName("VTK Cell Picker - 点击体元素进行高亮");
    renderWindow->SetSize(800, 600);
    
    // 创建交互器
    vtkSmartPointer<vtkRenderWindowInteractor> renderWindowInteractor = 
        vtkSmartPointer<vtkRenderWindowInteractor>::New();
    renderWindowInteractor->SetRenderWindow(renderWindow);
    
    // 创建自定义交互器样式
    vtkSmartPointer<MouseInteractorStyle> style = vtkSmartPointer<MouseInteractorStyle>::New();
    style->SetData(ugrid);
    style->SetActor(actor);
    renderWindowInteractor->SetInteractorStyle(style);
    
    // 开始渲染
    renderWindow->Render();
    
    // 输出使用说明
    std::cout << "===========================================" << std::endl;
    std::cout << "VTK 体网格单击拾取高亮演示程序" << std::endl;
    std::cout << "===========================================" << std::endl;
    std::cout << "操作说明:" << std::endl;
    std::cout << "- 左键单击网格上的体元素进行拾取和高亮" << std::endl;
    std::cout << "- 鼠标拖拽可以旋转视图" << std::endl;
    std::cout << "- 鼠标滚轮可以缩放" << std::endl;
    std::cout << "- 按 'q' 或关闭窗口退出程序" << std::endl;
    std::cout << "===========================================" << std::endl;
    
    // 启动交互
    renderWindowInteractor->Start();
    
    return 0;
}

// 编译说明:
// 需要链接以下VTK库:
// - vtkCommonCore
// - vtkCommonDataModel
// - vtkFiltersSources
// - vtkRenderingCore
// - vtkRenderingOpenGL2
// - vtkInteractionStyle
// - vtkFiltersExtraction
//
// CMake示例:
// find_package(VTK REQUIRED)
// target_link_libraries(your_target ${VTK_LIBRARIES})